{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building Code Documentation Agents with CrewAI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from crewai import LLM\n",
    "\n",
    "def load_llm():\n",
    "    llm = LLM(\n",
    "        # model=\"ollama/deepseek-r1:7b\",\n",
    "        model=\"ollama/llama3.2\",\n",
    "        base_url=\"http://localhost:11434\"\n",
    "    )\n",
    "    return llm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialization and Setup\n",
    "Initial imports for the CrewAI Flow and Crew and setting up the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing necessary libraries\n",
    "import yaml\n",
    "import subprocess\n",
    "from pathlib import Path\n",
    "from pydantic import BaseModel\n",
    "\n",
    "# Importing Crew related components\n",
    "from crewai import Agent, Task, Crew\n",
    "\n",
    "# Importing CrewAI Flow related components\n",
    "from crewai.flow.flow import Flow, listen, start\n",
    "\n",
    "# Apply a patch to allow nested asyncio loops in Jupyter\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the project URL\n",
    "\n",
    "In this demo, a sample repository is provided for you. However, feel free to test this on other public repositories! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "project_url = \"https://github.com/crewAIInc/nvidia-demo\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plan for our Flow\n",
    "\n",
    "1. Clone the repository for the project\n",
    "2. Plan the documentation for the project **[Planning Crew]** \n",
    "3. Create the documentation for the project **[Documentation Crew]**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Pydantic Schema\n",
    "\n",
    "Initial strucutre data we will use to capture the output of the planning crew"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define data structures to capture documentation planning output\n",
    "class DocItem(BaseModel):\n",
    "    \"\"\"Represents a documentation item\"\"\"\n",
    "    title: str\n",
    "    description: str\n",
    "    prerequisites: str\n",
    "    examples: list[str]\n",
    "    goal: str\n",
    "\n",
    "class DocPlan(BaseModel):\n",
    "    \"\"\"Documentation plan\"\"\"\n",
    "    overview: str\n",
    "    docs: list[DocItem]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optimizing for Llama 3.2 Prompting Template\n",
    "\n",
    "When using different models the ability to go a lower level and change the prompting template can drastically improve the performance of the model, you want to make sure to watch for the model's training prompt patterns and adjust accordingly.\n",
    "\n",
    "For Meta's Llama you can find it [in here](https://www.llama.com/docs/model-cards-and-prompt-formats/llama3_1/#prompt-template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Agents Prompting Template for Llama 3.3\n",
    "system_template=\"\"\"<|begin_of_text|><|start_header_id|>system<|end_header_id|>{{ .System }}<|eot_id|>\"\"\"\n",
    "prompt_template=\"\"\"<|start_header_id|>user<|end_header_id|>{{ .Prompt }}<|eot_id|>\"\"\"\n",
    "response_template=\"\"\"<|start_header_id|>assistant<|end_header_id|>{{ .Response }}<|eot_id|>\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Planning Crew\n",
    "\n",
    "Crew of AI Agents to strategize and create a documentation plan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from crewai_tools import (\n",
    "    DirectoryReadTool,\n",
    "    FileReadTool,\n",
    ")\n",
    "\n",
    "# Load agent and task configurations from YAML files\n",
    "with open('config/planner_agents.yaml', 'r') as f:\n",
    "    agents_config = yaml.safe_load(f)\n",
    "\n",
    "with open('config/planner_tasks.yaml', 'r') as f:\n",
    "    tasks_config = yaml.safe_load(f)\n",
    "\n",
    "code_explorer = Agent(\n",
    "  config=agents_config['code_explorer'],\n",
    "  system_template=system_template,\n",
    "  prompt_template=prompt_template,\n",
    "  response_template=response_template,\n",
    "  tools=[\n",
    "    DirectoryReadTool(),\n",
    "    FileReadTool()\n",
    "  ],\n",
    "  llm=load_llm()\n",
    ")\n",
    "documentation_planner = Agent(\n",
    "  config=agents_config['documentation_planner'],\n",
    "  system_template=system_template,\n",
    "  prompt_template=prompt_template,\n",
    "  response_template=response_template,\n",
    "  tools=[\n",
    "    DirectoryReadTool(),\n",
    "    FileReadTool()\n",
    "  ],\n",
    "  llm=load_llm()\n",
    ")\n",
    "\n",
    "analyze_codebase = Task(\n",
    "  config=tasks_config['analyze_codebase'],\n",
    "  agent=code_explorer\n",
    ")\n",
    "create_documentation_plan = Task(\n",
    "  config=tasks_config['create_documentation_plan'],\n",
    "  agent=documentation_planner,\n",
    "  output_pydantic=DocPlan\n",
    ")\n",
    "\n",
    "planning_crew = Crew(\n",
    "    agents=[code_explorer, documentation_planner],\n",
    "    tasks=[analyze_codebase, create_documentation_plan],\n",
    "    verbose=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Documentation Crew\n",
    "\n",
    "Crew of AI Agents to execute the documentation plan and create the documentation.\n",
    "Creating a guardrail to check the mermaid syntax in the documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "from crewai.tasks import TaskOutput\n",
    "import re\n",
    "\n",
    "def check_mermaid_syntax(task_output: TaskOutput):\n",
    "    text = task_output.raw\n",
    "\n",
    "    # Find all mermaid code blocks in the text\n",
    "    mermaid_blocks = re.findall(r'```mermaid\\n(.*?)\\n```', text, re.DOTALL)\n",
    "\n",
    "    for block in mermaid_blocks:\n",
    "        diagram_text = block.strip()\n",
    "        lines = diagram_text.split('\\n')\n",
    "        corrected_lines = []\n",
    "\n",
    "        for line in lines:\n",
    "            corrected_line = re.sub(r'\\|.*?\\|>', lambda match: match.group(0).replace('|>', '|'), line)\n",
    "            corrected_lines.append(corrected_line)\n",
    "\n",
    "        text = text.replace(block, \"\\n\".join(corrected_lines))\n",
    "\n",
    "    task_output.raw = text\n",
    "    return (True, task_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/akshay/Eigen/ai-engineering-hub/documentation-writer-flow/.venv/lib/python3.12/site-packages/ollama/_types.py:81: PydanticDeprecatedSince211: Accessing the 'model_fields' attribute on the instance is deprecated. Instead, you should access this attribute from the model class. Deprecated in Pydantic V2.11 to be removed in V3.0.\n",
      "  if key in self.model_fields:\n",
      "/Users/akshay/Eigen/ai-engineering-hub/documentation-writer-flow/.venv/lib/python3.12/site-packages/embedchain/embedder/ollama.py:27: LangChainDeprecationWarning: The class `OllamaEmbeddings` was deprecated in LangChain 0.3.1 and will be removed in 1.0.0. An updated version of the class exists in the :class:`~langchain-ollama package and should be used instead. To use it run `pip install -U :class:`~langchain-ollama` and import as `from :class:`~langchain_ollama import OllamaEmbeddings``.\n",
      "  embeddings = OllamaEmbeddings(model=self.config.model, base_url=config.base_url)\n"
     ]
    }
   ],
   "source": [
    "from crewai_tools import (\n",
    "    DirectoryReadTool,\n",
    "    FileReadTool,\n",
    "    WebsiteSearchTool\n",
    ")\n",
    "\n",
    "# Load agent and task configurations from YAML files\n",
    "with open('config/documentation_agents.yaml', 'r') as f:\n",
    "    agents_config = yaml.safe_load(f)\n",
    "\n",
    "with open('config/documentation_tasks.yaml', 'r') as f:\n",
    "    tasks_config = yaml.safe_load(f)\n",
    "\n",
    "overview_writer = Agent(config=agents_config['overview_writer'], tools=[\n",
    "    DirectoryReadTool(),\n",
    "    FileReadTool(),\n",
    "    WebsiteSearchTool(\n",
    "      website=\"https://mermaid.js.org/intro/\",\n",
    "      config=dict(\n",
    "        embedder=dict(\n",
    "            provider=\"ollama\",\n",
    "            config=dict(\n",
    "                model=\"nomic-embed-text\",\n",
    "            ),\n",
    "        )\n",
    "        )\n",
    "      )\n",
    "  ],\n",
    "  llm=load_llm()\n",
    ")\n",
    "\n",
    "documentation_reviewer = Agent(config=agents_config['documentation_reviewer'], tools=[\n",
    "    DirectoryReadTool(directory=\"docs/\", name=\"Check existing documentation folder\"),\n",
    "    FileReadTool(),\n",
    "  ],\n",
    "  llm=load_llm()\n",
    ")\n",
    "\n",
    "draft_documentation = Task(\n",
    "  config=tasks_config['draft_documentation'],\n",
    "  agent=overview_writer\n",
    ")\n",
    "\n",
    "qa_review_documentation = Task(\n",
    "  config=tasks_config['qa_review_documentation'],\n",
    "  agent=documentation_reviewer,\n",
    "  guardrail=check_mermaid_syntax,\n",
    "  max_retries=5\n",
    ")\n",
    "\n",
    "documentation_crew = Crew(\n",
    "    agents=[overview_writer, documentation_reviewer],\n",
    "    tasks=[draft_documentation, qa_review_documentation],\n",
    "    verbose=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Documentation Flow\n",
    "\n",
    "A Flow to create the documentation for the project where we will use the planning crew to plan the documentation and the documentation crew to create the documentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from typing import List\n",
    "\n",
    "\n",
    "class DocumentationState(BaseModel):\n",
    "  \"\"\"\n",
    "  State for the documentation flow\n",
    "  \"\"\"\n",
    "  project_url: str = project_url\n",
    "  repo_path: Path = \"workdir/\"\n",
    "  docs: List[str] = []\n",
    "\n",
    "class CreateDocumentationFlow(Flow[DocumentationState]):\n",
    "  # Clone the repository, initial step\n",
    "  # No need for AI Agents on this step, so we just use regular Python code\n",
    "  @start()\n",
    "  def clone_repo(self):\n",
    "    print(f\"# Cloning repository: {self.state.project_url}\\n\")\n",
    "    # Extract repo name from URL\n",
    "    repo_name = self.state.project_url.split(\"/\")[-1]\n",
    "    self.state.repo_path = f\"{self.state.repo_path}{repo_name}\"\n",
    "\n",
    "  # Check if directory exists\n",
    "    if Path(self.state.repo_path).exists():\n",
    "      print(f\"# Repository directory already exists at {self.state.repo_path}\\n\")\n",
    "      subprocess.run([\"rm\", \"-rf\", self.state.repo_path])\n",
    "      print(\"# Removed existing directory\\n\")\n",
    "\n",
    "    # Clone the repository\n",
    "    subprocess.run([\"git\", \"clone\", self.state.project_url, self.state.repo_path])\n",
    "    return self.state\n",
    "\n",
    "  @listen(clone_repo)\n",
    "  def plan_docs(self):\n",
    "    print(f\"# Planning documentation for: {self.state.repo_path}\\n\")\n",
    "    result = planning_crew.kickoff(inputs={'repo_path': self.state.repo_path})\n",
    "    print(f\"# Planned docs for {self.state.repo_path}:\")\n",
    "    for doc in result.pydantic.docs:\n",
    "        print(f\"    - {doc.title}\")\n",
    "    return result\n",
    "\n",
    "  @listen(plan_docs)\n",
    "  def save_plan(self, plan):\n",
    "    with open(\"docs/plan.json\", \"w\") as f:\n",
    "      f.write(plan.raw)\n",
    "\n",
    "  @listen(plan_docs)\n",
    "  def create_docs(self, plan):\n",
    "    for doc in plan.pydantic.docs:\n",
    "      print(f\"\\n# Creating documentation for: {doc.title}\")\n",
    "      result = documentation_crew.kickoff(inputs={\n",
    "        'repo_path': self.state.repo_path,\n",
    "        'title': doc.title,\n",
    "        'overview': plan.pydantic.overview,\n",
    "        'description': doc.description,\n",
    "        'prerequisites': doc.prerequisites,\n",
    "        'examples': '\\n'.join(doc.examples),\n",
    "        'goal': doc.goal\n",
    "      })\n",
    "\n",
    "      # Save documentation to file in docs folder\n",
    "      docs_dir = Path(\"docs\")\n",
    "      docs_dir.mkdir(exist_ok=True)\n",
    "      title = doc.title.lower().replace(\" \", \"_\") + \".mdx\"\n",
    "      self.state.docs.append(str(docs_dir / title))\n",
    "      with open(docs_dir / title, \"w\") as f:\n",
    "          f.write(result.raw)\n",
    "    print(f\"\\n# Documentation created for: {self.state.repo_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Documentation Flow\n",
    "\n",
    "After running this cell, check the `docs` directory for the generated documentation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #000080; text-decoration-color: #000080\">╭──────────────────────────────────────────────── Flow Execution ─────────────────────────────────────────────────╮</span>\n",
       "<span style=\"color: #000080; text-decoration-color: #000080\">│</span>                                                                                                                 <span style=\"color: #000080; text-decoration-color: #000080\">│</span>\n",
       "<span style=\"color: #000080; text-decoration-color: #000080\">│</span>  <span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">Starting Flow Execution</span>                                                                                        <span style=\"color: #000080; text-decoration-color: #000080\">│</span>\n",
       "<span style=\"color: #000080; text-decoration-color: #000080\">│</span>  <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">Name: </span><span style=\"color: #000080; text-decoration-color: #000080\">CreateDocumentationFlow</span>                                                                                  <span style=\"color: #000080; text-decoration-color: #000080\">│</span>\n",
       "<span style=\"color: #000080; text-decoration-color: #000080\">│</span>  <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">ID: </span><span style=\"color: #000080; text-decoration-color: #000080\">7e73bcbc-5fd6-4742-97f7-2f28264632af</span>                                                                       <span style=\"color: #000080; text-decoration-color: #000080\">│</span>\n",
       "<span style=\"color: #000080; text-decoration-color: #000080\">│</span>                                                                                                                 <span style=\"color: #000080; text-decoration-color: #000080\">│</span>\n",
       "<span style=\"color: #000080; text-decoration-color: #000080\">│</span>                                                                                                                 <span style=\"color: #000080; text-decoration-color: #000080\">│</span>\n",
       "<span style=\"color: #000080; text-decoration-color: #000080\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[34m╭─\u001b[0m\u001b[34m───────────────────────────────────────────────\u001b[0m\u001b[34m Flow Execution \u001b[0m\u001b[34m────────────────────────────────────────────────\u001b[0m\u001b[34m─╮\u001b[0m\n",
       "\u001b[34m│\u001b[0m                                                                                                                 \u001b[34m│\u001b[0m\n",
       "\u001b[34m│\u001b[0m  \u001b[1;34mStarting Flow Execution\u001b[0m                                                                                        \u001b[34m│\u001b[0m\n",
       "\u001b[34m│\u001b[0m  \u001b[37mName: \u001b[0m\u001b[34mCreateDocumentationFlow\u001b[0m                                                                                  \u001b[34m│\u001b[0m\n",
       "\u001b[34m│\u001b[0m  \u001b[37mID: \u001b[0m\u001b[34m7e73bcbc-5fd6-4742-97f7-2f28264632af\u001b[0m                                                                       \u001b[34m│\u001b[0m\n",
       "\u001b[34m│\u001b[0m                                                                                                                 \u001b[34m│\u001b[0m\n",
       "\u001b[34m│\u001b[0m                                                                                                                 \u001b[34m│\u001b[0m\n",
       "\u001b[34m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">🌊 Flow: </span><span style=\"color: #000080; text-decoration-color: #000080\">CreateDocumentationFlow</span>\n",
       "<span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">    ID: </span><span style=\"color: #000080; text-decoration-color: #000080\">7e73bcbc-5fd6-4742-97f7-2f28264632af</span>\n",
       "└── <span style=\"color: #808000; text-decoration-color: #808000\">🧠 Starting Flow...</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;34m🌊 Flow: \u001b[0m\u001b[34mCreateDocumentationFlow\u001b[0m\n",
       "\u001b[37m    ID: \u001b[0m\u001b[34m7e73bcbc-5fd6-4742-97f7-2f28264632af\u001b[0m\n",
       "└── \u001b[33m🧠 Starting Flow...\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[35m Flow started with ID: 7e73bcbc-5fd6-4742-97f7-2f28264632af\u001b[00m\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">🌊 Flow: </span><span style=\"color: #000080; text-decoration-color: #000080\">CreateDocumentationFlow</span>\n",
       "<span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">    ID: </span><span style=\"color: #000080; text-decoration-color: #000080\">7e73bcbc-5fd6-4742-97f7-2f28264632af</span>\n",
       "├── <span style=\"color: #808000; text-decoration-color: #808000\">🧠 Starting Flow...</span>\n",
       "└── <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">🔄 Running: clone_repo</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;34m🌊 Flow: \u001b[0m\u001b[34mCreateDocumentationFlow\u001b[0m\n",
       "\u001b[37m    ID: \u001b[0m\u001b[34m7e73bcbc-5fd6-4742-97f7-2f28264632af\u001b[0m\n",
       "├── \u001b[33m🧠 Starting Flow...\u001b[0m\n",
       "└── \u001b[1;33m🔄 Running:\u001b[0m\u001b[1;33m clone_repo\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Cloning repository: https://github.com/crewAIInc/nvidia-demo\n",
      "\n",
      "# Repository directory already exists at workdir/nvidia-demo\n",
      "\n",
      "# Removed existing directory\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cloning into 'workdir/nvidia-demo'...\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">🌊 Flow: </span><span style=\"color: #000080; text-decoration-color: #000080\">CreateDocumentationFlow</span>\n",
       "<span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">    ID: </span><span style=\"color: #000080; text-decoration-color: #000080\">7e73bcbc-5fd6-4742-97f7-2f28264632af</span>\n",
       "├── <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">Flow Method Step</span>\n",
       "└── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: clone_repo</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;34m🌊 Flow: \u001b[0m\u001b[34mCreateDocumentationFlow\u001b[0m\n",
       "\u001b[37m    ID: \u001b[0m\u001b[34m7e73bcbc-5fd6-4742-97f7-2f28264632af\u001b[0m\n",
       "├── \u001b[37mFlow Method Step\u001b[0m\n",
       "└── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m clone_repo\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">🌊 Flow: </span><span style=\"color: #000080; text-decoration-color: #000080\">CreateDocumentationFlow</span>\n",
       "<span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">    ID: </span><span style=\"color: #000080; text-decoration-color: #000080\">7e73bcbc-5fd6-4742-97f7-2f28264632af</span>\n",
       "├── <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">Flow Method Step</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: clone_repo</span>\n",
       "└── <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">🔄 Running: plan_docs</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;34m🌊 Flow: \u001b[0m\u001b[34mCreateDocumentationFlow\u001b[0m\n",
       "\u001b[37m    ID: \u001b[0m\u001b[34m7e73bcbc-5fd6-4742-97f7-2f28264632af\u001b[0m\n",
       "├── \u001b[37mFlow Method Step\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m clone_repo\u001b[0m\n",
       "└── \u001b[1;33m🔄 Running:\u001b[0m\u001b[1;33m plan_docs\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Planning documentation for: workdir/nvidia-demo\n",
      "\n",
      "# Planned docs for workdir/nvidia-demo:\n",
      "    - Technical Overview\n",
      "    - Component Breakdown\n",
      "    - CUDA Shared Libraries\n",
      "    - Design Patterns\n",
      "    - API Documentation\n",
      "    - Data Flow\n",
      "    - Design Considerations and Best Practices\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">🌊 Flow: </span><span style=\"color: #000080; text-decoration-color: #000080\">CreateDocumentationFlow</span>\n",
       "<span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">    ID: </span><span style=\"color: #000080; text-decoration-color: #000080\">7e73bcbc-5fd6-4742-97f7-2f28264632af</span>\n",
       "├── <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">Flow Method Step</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: clone_repo</span>\n",
       "└── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: plan_docs</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;34m🌊 Flow: \u001b[0m\u001b[34mCreateDocumentationFlow\u001b[0m\n",
       "\u001b[37m    ID: \u001b[0m\u001b[34m7e73bcbc-5fd6-4742-97f7-2f28264632af\u001b[0m\n",
       "├── \u001b[37mFlow Method Step\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m clone_repo\u001b[0m\n",
       "└── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m plan_docs\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">🌊 Flow: </span><span style=\"color: #000080; text-decoration-color: #000080\">CreateDocumentationFlow</span>\n",
       "<span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">    ID: </span><span style=\"color: #000080; text-decoration-color: #000080\">7e73bcbc-5fd6-4742-97f7-2f28264632af</span>\n",
       "├── <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">Flow Method Step</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: clone_repo</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: plan_docs</span>\n",
       "└── <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">🔄 Running: save_plan</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;34m🌊 Flow: \u001b[0m\u001b[34mCreateDocumentationFlow\u001b[0m\n",
       "\u001b[37m    ID: \u001b[0m\u001b[34m7e73bcbc-5fd6-4742-97f7-2f28264632af\u001b[0m\n",
       "├── \u001b[37mFlow Method Step\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m clone_repo\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m plan_docs\u001b[0m\n",
       "└── \u001b[1;33m🔄 Running:\u001b[0m\u001b[1;33m save_plan\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">🌊 Flow: </span><span style=\"color: #000080; text-decoration-color: #000080\">CreateDocumentationFlow</span>\n",
       "<span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">    ID: </span><span style=\"color: #000080; text-decoration-color: #000080\">7e73bcbc-5fd6-4742-97f7-2f28264632af</span>\n",
       "├── <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">Flow Method Step</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: clone_repo</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: plan_docs</span>\n",
       "└── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: save_plan</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;34m🌊 Flow: \u001b[0m\u001b[34mCreateDocumentationFlow\u001b[0m\n",
       "\u001b[37m    ID: \u001b[0m\u001b[34m7e73bcbc-5fd6-4742-97f7-2f28264632af\u001b[0m\n",
       "├── \u001b[37mFlow Method Step\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m clone_repo\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m plan_docs\u001b[0m\n",
       "└── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m save_plan\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">🌊 Flow: </span><span style=\"color: #000080; text-decoration-color: #000080\">CreateDocumentationFlow</span>\n",
       "<span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">    ID: </span><span style=\"color: #000080; text-decoration-color: #000080\">7e73bcbc-5fd6-4742-97f7-2f28264632af</span>\n",
       "├── <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">Flow Method Step</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: clone_repo</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: plan_docs</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: save_plan</span>\n",
       "└── <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">🔄 Running: create_docs</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;34m🌊 Flow: \u001b[0m\u001b[34mCreateDocumentationFlow\u001b[0m\n",
       "\u001b[37m    ID: \u001b[0m\u001b[34m7e73bcbc-5fd6-4742-97f7-2f28264632af\u001b[0m\n",
       "├── \u001b[37mFlow Method Step\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m clone_repo\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m plan_docs\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m save_plan\u001b[0m\n",
       "└── \u001b[1;33m🔄 Running:\u001b[0m\u001b[1;33m create_docs\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "# Creating documentation for: Technical Overview\n",
      "\n",
      "# Creating documentation for: Component Breakdown\n",
      "\n",
      "# Creating documentation for: CUDA Shared Libraries\n",
      "\n",
      "# Creating documentation for: Design Patterns\n",
      "\n",
      "# Creating documentation for: API Documentation\n",
      "\n",
      "# Creating documentation for: Data Flow\n",
      "\n",
      "# Creating documentation for: Design Considerations and Best Practices\n",
      "\n",
      "# Documentation created for: workdir/nvidia-demo\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">🌊 Flow: </span><span style=\"color: #000080; text-decoration-color: #000080\">CreateDocumentationFlow</span>\n",
       "<span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">    ID: </span><span style=\"color: #000080; text-decoration-color: #000080\">7e73bcbc-5fd6-4742-97f7-2f28264632af</span>\n",
       "├── <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">Flow Method Step</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: clone_repo</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: plan_docs</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: save_plan</span>\n",
       "└── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: create_docs</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;34m🌊 Flow: \u001b[0m\u001b[34mCreateDocumentationFlow\u001b[0m\n",
       "\u001b[37m    ID: \u001b[0m\u001b[34m7e73bcbc-5fd6-4742-97f7-2f28264632af\u001b[0m\n",
       "├── \u001b[37mFlow Method Step\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m clone_repo\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m plan_docs\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m save_plan\u001b[0m\n",
       "└── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m create_docs\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Flow Finished: </span><span style=\"color: #008000; text-decoration-color: #008000\">CreateDocumentationFlow</span>\n",
       "├── <span style=\"color: #c0c0c0; text-decoration-color: #c0c0c0\">Flow Method Step</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: clone_repo</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: plan_docs</span>\n",
       "├── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: save_plan</span>\n",
       "└── <span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">✅ Completed: create_docs</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;32m✅ Flow Finished: \u001b[0m\u001b[32mCreateDocumentationFlow\u001b[0m\n",
       "├── \u001b[37mFlow Method Step\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m clone_repo\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m plan_docs\u001b[0m\n",
       "├── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m save_plan\u001b[0m\n",
       "└── \u001b[1;32m✅ Completed:\u001b[0m\u001b[1;32m create_docs\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "flow = CreateDocumentationFlow()\n",
    "flow.kickoff()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot One of the Documents\n",
    "\n",
    "Let's visualize one of the generated documentation files to verify the output. This will help us ensure the documentation was created successfully and formatted correctly.\n",
    "\n",
    "The generated documentation files can be found in the `docs` directory in the root of the project. Each documentation file is saved with a `.mdx` extension and follows the naming convention of lowercase words separated by underscores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Documentation files generated:\n",
      "- docs/core_workflows_and_data_flows.mdx\n",
      "- docs/technical_overview.mdx\n",
      "- docs/component_breakdown.mdx\n",
      "- docs/design_patterns.mdx\n",
      "- docs/getting_started_guide.mdx\n",
      "- docs/data_flow.mdx\n",
      "- docs/api_documentation.mdx\n",
      "- docs/project_overview_and_architecture.mdx\n",
      "- docs/quality_assurance_in_documentation.mdx\n",
      "- docs/design_considerations_and_best_practices.mdx\n",
      "- docs/comprehensive_documentation_strategy.mdx\n",
      "- docs/cuda_shared_libraries.mdx\n",
      "\n",
      "Displaying contents of first doc:\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "<think>\n",
       "Okay, I'm trying to help validate the documentation for workdir/nvidia-demo. The user has given me a detailed task with several criteria to follow. Let me break this down step by step.\n",
       "\n",
       "First, I need to check if all the technical accuracy aspects are covered. That means ensuring that every architectural description in the docs matches the actual code, checking component relationships and interactions, validating code examples with tests and usage, and confirming that mermaid diagrams reflect real data flows.\n",
       "\n",
       "Next, for documentation completeness, I have to verify that all key components are documented, ensure existing workflows in code are covered, check integration patterns, and confirm that troubleshooting scenarios are accurate.\n",
       "\n",
       "Then, looking at the quality part, I need to remove any speculative or unimplemented features, update examples to match current code, make sure mermaid diagrams enhance understanding, not wrap them in fences or meta-comments, and keep it clean without images or media files.\n",
       "\n",
       "Now, checking the context provided: The project uses CUDA with NVIDIA libraries like nv Hardy and cusolver. Setup includes cloning repo, installing dependencies, initializing CUDA contexts, setting env variables.\n",
       "\n",
       "Components include CUDA Kernel Development, GPU Resource Management, Performance Analysis. Each has examples in code snippets.\n",
       "\n",
       "High-Level Flow diagram is present but not described here. There are two Mermaid diagrams: one showing component relationships and another data flow process. The code examples provided seem accurate but maybe need updating if new functions are added or old ones deprecated.\n",
       "\n",
       "I should use the tools to list content, check each section for consistency with codebase, update examples as needed, ensure all components are covered, and validate that mermaid diagrams match actual flows without extra fluff.\n",
       "\n",
       "I'll start by checking existing docs using the folder tool. Then read each file's content, especially the setup, components, flow, and examples sections. I'll cross-reference them with code to spot any discrepancies or missing parts. Finally, make sure all criteria are met before finalizing.\n",
       "</think>\n",
       "\n",
       "Thought: I have reviewed the documentation against the project setup and components, ensuring alignment with the codebase.\n",
       "\n",
       "Action:\n",
       "- Check existing documentation folder\n",
       "- Read a file's content (for setup instructions)\n",
       "- Read a file's content (for CUDA Kernel Development example)\n",
       "- Read a file's content (for GPU Resource Management example)\n",
       "- Read a file's content (for Performance Analysis example)\n",
       "- Read a file's content (for High-Level Flow diagram)\n",
       "- Read a file's content (for Component Relationships Mermaid diagram)\n",
       "- Read a file's content (for Data Flow Process Mermaid diagram)\n"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# List all files in docs folder and display the first doc using IPython.display\n",
    "from IPython.display import Markdown\n",
    "import pathlib\n",
    "\n",
    "docs_dir = pathlib.Path(\"docs\")\n",
    "print(\"Documentation files generated:\")\n",
    "for doc_file in docs_dir.glob(\"*.mdx\"):\n",
    "    print(f\"- docs/{doc_file.name}\")\n",
    "\n",
    "print(\"\\nDisplaying contents of first doc:\\n\")\n",
    "first_doc = pathlib.Path(flow.state.docs[0]).read_text()\n",
    "display(Markdown(first_doc))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
